"""
Variational Data Assimilation Framework (VDAF)
Complete implementation with examples for both 3D-Var and 4D-Var
"""
import numpy as np
from scipy.optimize import minimize
from typing import Dict, List, Optional, Tuple
import matplotlib.pyplot as plt
import random
from Tools.AssimilationSystem import *
from Tools.var_methods import *
from Tools.adapters import *
from Tools.var4d_capture import create_var4d_with_capture
from Tools.barotropic_model import BarotropicModel
from Tools.barotropic_adapter import BarotropicObservationOperator
from Tools.model_factory import ModelFactory
import time

def run_3dvar(x_b, x_truth, model, state_size, obs_size, all_obs_ops, all_observations):
    """Run 3D-Var assimilation with t=0 observations only."""
    return
    print("\n=== Running 3D-Var Example ===")
    config_3dvar = AssimilationConfig(
        method='3DVar',
        state_size=state_size,
        obs_size=obs_size,
        time_window=0,
        max_outer_loops=1,
        max_inner_loops=100,
        bg_error_variance=1.0,  # Increased background variance
        obs_error_variance=0.1   # Increased obs error variance
    )
    
    # Create 3D-Var system with t=0 observations only
    observations_3dvar = {0: all_observations[0]}
    assim_system_3dvar = AssimilationSystem(config_3dvar, model)
    assim_system_3dvar.set_observation_operators([all_obs_ops[0]])
    x_a_3dvar = assim_system_3dvar.assimilate(x_b, observations_3dvar)
    
    # Calculate 3D-Var error statistics
    rmse_b = np.sqrt(np.mean((x_b - x_truth)**2))
    rmse_a_3dvar = np.sqrt(np.mean((x_a_3dvar - x_truth)**2))
    
    print(f"\nResults for 3D-Var:")
    print(f"Background RMSE: {rmse_b:.4f}")
    print(f"Analysis RMSE: {rmse_a_3dvar:.4f}")
    print(f"Improvement: {(1 - rmse_a_3dvar/rmse_b)*100:.1f}%")
    
    # Visualize 3D-Var results
    # visualize_results(x_truth, x_b, x_a_3dvar, observations_3dvar, [all_obs_ops[0]], 
    #                  "3D-Var Assimilation Results")
    
    return x_a_3dvar, rmse_b, rmse_a_3dvar


def plot_state_comparison(model, x_truth, x_b, x_a, title_suffix="", filename="comparison_results.png"):
    """Plot comparison between truth, background and analysis state."""
    plt.figure(figsize=(15, 5))
    
    # Plot vorticity fields
    plt.subplot(1, 3, 1)
    plt.title(f"Truth {title_suffix}")
    truth_2d = model._reshape_state(x_truth)
    plt.imshow(truth_2d, cmap='RdBu_r')
    plt.colorbar(label='Vorticity')
    
    plt.subplot(1, 3, 2)
    plt.title(f"Background {title_suffix}")
    bg_2d = model._reshape_state(x_b)
    plt.imshow(bg_2d, cmap='RdBu_r')
    plt.colorbar(label='Vorticity')
    
    plt.subplot(1, 3, 3)
    plt.title(f"Analysis {title_suffix}")
    analysis_2d = model._reshape_state(x_a)
    plt.imshow(analysis_2d, cmap='RdBu_r')
    plt.colorbar(label='Vorticity')
    
    plt.tight_layout()
    plt.savefig(filename)
    plt.close()

def run_drp4dvar(x_b, x_truth, model, state_size, obs_size, max_time_window, bg_error_variance, 
                obs_error_variance, all_obs_ops, all_observations):
    """Run DRP4DVar assimilation with observations from all time steps."""
    print("\n=== Running DRP4DVar Example ===")
    config_4dvar = AssimilationConfig(
        method='DRP4DVar',
        state_size=state_size,
        obs_size=obs_size,
        time_window=max_time_window,
        max_outer_loops=1,
        max_inner_loops=50,
        bg_error_variance=bg_error_variance,
        obs_error_variance=obs_error_variance
    )
    config_4dvar.ensemble_size = 60
    
    # Run DRP4DVar with the same initial conditions
    assim_system_4dvar = AssimilationSystem(config_4dvar, model)
    assim_system_4dvar.set_observation_operators(all_obs_ops)
    x_a_4dvar = assim_system_4dvar.assimilate(x_b, all_observations)
    
    # Calculate DRP4DVar error statistics
    rmse_b = np.sqrt(np.mean((x_b - x_truth)**2))
    rmse_a_4dvar = np.sqrt(np.mean((x_a_4dvar - x_truth)**2))
    
    print(f"\nResults for DRP4DVar:")
    print(f"max inner loops: {config_4dvar.max_inner_loops}")
    #print(f"ensemble size: {config_4dvar.ensemble_size}")
    print(f"Background RMSE: {rmse_b:.4f}")
    print(f"Analysis RMSE: {rmse_a_4dvar:.4f}")
    print(f"Improvement: {(1 - rmse_a_4dvar/rmse_b)*100:.1f}%")
    # plot_state_comparison(model, x_truth, x_b, x_a_4dvar, "", "barotropic_drp4dvar_results.png")
    return x_a_4dvar, rmse_b, rmse_a_4dvar

# Main execution
if __name__ == "__main__":
    # Set fixed seeds for reproducibility
    RANDOM_SEED = 42
    random.seed(RANDOM_SEED)
    np.random.seed(RANDOM_SEED)
    
    # Create configuration for barotropic model
    state_size = 300  # Needs to be composite number for grid dimensions
    obs_size = state_size // 4     # 25% of state size
    bg_error_variance = 0.1
    obs_error_variance = 0.05 
    
    # Create barotropic model
    model = ModelFactory.create_model("BAROTROPIC", state_size)
    
    # Get actual state size from model (important for barotropic grid dimensions)
    state_size = model.state_size
    obs_size = state_size // 4  # 25% of state points
    
    # Create barotropic observation operators for all time steps
    max_time_window = 4  # Maximum time window for 4D-Var
    all_obs_ops = [
        BarotropicObservationOperator(
            model, 
            state_size,
            obs_size,
            obs_type='vorticity',
            seed=42 + t
        ) for t in range(max_time_window + 1)
    ]
    
    # Generate one set of synthetic data with the 4D-Var time window
    config_data = AssimilationConfig(
        method='DRP4DVar',
        state_size=state_size,
        obs_size=obs_size,
        time_window=max_time_window,
        bg_error_variance=bg_error_variance,
        obs_error_variance=obs_error_variance
    )

    # Generate synthetic data with barotropic model
    x_truth = model.generate_random_state(mean_amplitude=1.0, seed=RANDOM_SEED)
    
    # Generate background by perturbing truth
    error = np.random.multivariate_normal(np.zeros(state_size), config_data.b_matrix)
    x_b = x_truth + error
    start = time.perf_counter()
    # Generate observations by running truth model
    truth_traj = model.integrate(x_truth, max_time_window)
    elapsed_truth_traj = time.perf_counter() - start
    print(f"integrate time: {elapsed_truth_traj:.3f} seconds")
    start = time.perf_counter()
    all_observations = {}
    for t in range(max_time_window + 1):
        # Apply observation operator
        y_perfect = all_obs_ops[t].h(truth_traj[t])
        
        # Add observation error
        obs_error = np.random.multivariate_normal(
            np.zeros(obs_size), 
            config_data.r_matrix
        )
        y_obs = y_perfect + obs_error
        all_observations[t] = y_obs
    elapsed_observations = time.perf_counter() - start
    print(f"observations time: {elapsed_observations:.3f} seconds")
    # Run the three assimilation methods
    # start = time.perf_counter()
    # x_a_3dvar, rmse_b, rmse_a_3dvar = run_3dvar(x_b, x_truth, model, state_size, obs_size, 
    #                                            all_obs_ops, all_observations)
    # elapsed_3dvar = time.perf_counter() - start
    # print(f"run_3dvar time: {elapsed_3dvar:.3f} seconds")
    start = time.perf_counter()
    x_a_drp4dvar, _, rmse_a_drp4dvar = run_drp4dvar(x_b, x_truth, model, state_size, obs_size, 
                                                  max_time_window, bg_error_variance, 
                                                  obs_error_variance, all_obs_ops, all_observations)
    elapsed_drp4dvar = time.perf_counter() - start 
    print(f"run_drp4dvar time: {elapsed_drp4dvar:.3f} seconds")
    
    # Visualize 4D-Var results (showing t=0 observations for comparison)
    # observations_t0 = {0: all_observations[0]}
    # visualize_results(x_truth, x_b, x_a_4dvar, observations_t0, [all_obs_ops[0]], 
    #                  "4D-Var Assimilation Results")

    # start = time.perf_counter()
    # xx_adapter, yy_adapter, hmm_adapter, var4d_with_capture = run_var4d_from_var_method(
    #     x_b, x_truth, model, state_size, obs_size, max_time_window, 
    #     all_obs_ops, all_observations)
    # elapsed_var4d_from_var_method = time.perf_counter() - start
    # print(f"run_var4d_from_var_method time: {elapsed_var4d_from_var_method:.3f} seconds")
    
    # start = time.perf_counter()
    # x_a_var4d, rmse_a_var4d = evaluate_var4d_results(
    #      x_b, x_truth, observations_t0, 
    #     all_obs_ops, rmse_a_3dvar, rmse_a_3dvar, rmse_b
    # )
    # elapsed_evaluate_var4d_results = time.perf_counter() - start
    # print(f"evaluate_var4d_results time: {elapsed_evaluate_var4d_results:.3f} seconds")
